import os
import random
import torch
import torch.nn as nn
import torch.utils.data as Data     # 小批量
import torchvision                  # 提供图像数据集
import matplotlib.pyplot as plt     # 绘图
from matplotlib import cm
import utils
from cnn import CNN


EPOCH = 1               # 轮次
BATCH_SIZE = 50         # 批尺寸
LR = 0.001              # 学习率

MNIST_ROOT_DIR = 'e:/sfxData/DeepLearning'

cnn = CNN()
# print(cnn)
'''
CNN(
  (conv1): Sequential(
    (0): Conv2d(1, 16, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))
    (1): ReLU()
    (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  )
  (conv2): Sequential(
    (0): Conv2d(16, 32, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))
    (1): ReLU()
    (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  )
  (out): Linear(in_features=1568, out_features=10, bias=True)
)
'''

train_data = utils.get_mnist(MNIST_ROOT_DIR, train=True)
# utils.disp_mnist(train_data)
# utils.mnist_image(train_data, idx=random.randint(0, 60000))

# 得到小批量数据 (shuffle - 打乱顺序)
#
train_loader = Data.DataLoader(
    dataset=train_data, batch_size=BATCH_SIZE, shuffle=True)

# train_loader.dataset
for step, (b_x, b_y) in enumerate(train_loader):
	print(step)
	output = cnn(b_x)[0]
	print(output)


x = 0
